Greptile Summary

This PR enables head_dim=256 support for FlashAttention 4 by delegating head-dimension validation to FA4's own _validate_head_dims function, bumping the required FA4 version to 4.0.0b11, and adding a dedicated SM100-only test. A cross-attention fallback guard is also added for the hd256 path on SM100, where the dedicated kernel is known to diverge.

• Backend change: The hard-coded SM100/SM90 head-dim table in get_attention_backend is replaced with a live call to FlashAttentionUtils.v4_validate_head_dims, keeping TE in sync with FA4's supported shapes automatically.
• New test: test_dpa_fa4_hdim256 is added with explicit SM100/103 skipif guard (separate from test_dpa_fa4_base), and cuDNN version guards are removed from all FA4 tests since FA4 doesn't use cuDNN.
• Cross-attention fallback: FA4 is disabled on SM100 for head_dim=256 when max_seqlen_q != max_seqlen_kv, guarding against the dedicated hd256 kernel's known divergence on decode-like shapes.

Confidence Score: 4/5

Safe to merge on SM100 hardware; the static head-dim table replacement and dedicated hd256 test are well-structured, but the single FA4 import statement still hard-fails for any FA4 build that does not export _validate_head_dims.

The core logic change — delegating head-dim validation to FA4's own function — is correct and the SM100-gated test properly signals intent. The remaining concern is that _validate_head_dims is imported in the same from … import block as the two main FA4 symbols: any older FA4 build missing that private symbol raises an unhandled ImportError at module load time, breaking all FA4 paths for those users. That import structure was flagged in the prior review round and is still unchanged in this PR.

transformer_engine/pytorch/attention/dot_product_attention/backends.py — the grouped FA4 import at lines 167–171 is the single point where a missing _validate_head_dims symbol would crash module load for users on older FA4 builds.

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[get_attention_backend called] --> B{use_flash_attention_4\nAND v4_is_installed\nAND v4_validate_head_dims ≠ None?}
    B -- No --> Z[Skip FA4 head-dim validation]
    B -- Yes --> C[Compute _fa4_alignment\n= 16 // element_size]
    C --> D[Call v4_validate_head_dims\nhead_dim_qk, head_dim_v,\nsm_major, alignment]
    D -- AssertionError --> E[log: unsupported head dims\nuse_flash_attention_4 = False]
    D -- OK --> F{SM100 AND\nhd_qk == hd_v == 256\nAND max_seqlen_q != max_seqlen_kv?}
    F -- Yes --> G[log: hd256 cross-attn fallback\nuse_flash_attention_4 = False]
    F -- No --> H{is_training AND\nhd_qk != hd_v AND\nhd_qk >= 128 AND SM100?}
    H -- Yes --> I{dV TMEM misalignment?\ntile_hdimv//2 % dk_reduce_ncol != 0}
    I -- Yes --> J[log: MLA dV bug\nuse_flash_attention_4 = False]
    I -- No --> K[FA4 enabled]
    H -- No --> K
    E --> L[Fall back to other backend]
    G --> L
    J --> L

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into xiny/headdim256..." | Re-trigger Greptile}

/te-ci pytorch L3

@vcherepanov-nv @KshitijLakhani Please review.

LGTM

/te-ci pytorch L3

/te-ci pytorch L3